package com.example.learn

import java.util.*

fun main(){
    val root = TreeNode(1)
    val node1 = TreeNode(2)
    val node2 = TreeNode(3)
    val node3 = TreeNode(4)
    val node4 = TreeNode(5)
    val node5 = TreeNode(6)
    root.left = node1
    root.right = node2
    node1.left = node3
    node1.right = node4
    node2.right = node5
    print(averageOfLevels(root))
}

fun averageOfLevels(root: TreeNode) : List<Double> {
    val averages = mutableListOf<Double>()
    val queue = LinkedList<TreeNode>()
    queue.add(root)
    while (!queue.isEmpty()) {
        val size = queue.size
        var ave = 0.0
        for (i in 0 until size) {
            val node = queue.poll()
            if (node.left != null) queue.add(node.left!!)
            if (node.right != null) queue.add(node.right!!)
            ave += node.value
        }
        ave = ave * 1.0 / size
        averages.add(ave)
    }
    return averages
}

class TreeNode(val value: Int, var left: TreeNode? = null, var right: TreeNode? = null)